#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

#include <unistd.h>
#include <poll.h>

#include <linux/types.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>
//#include <netpacket/packet.h>
#include <net/ethernet.h>
#include <linux/if_ether.h>
#include <linux/if_arp.h>
#include <ifaddrs.h>

// ================================================================================

#define ET_IPV4 0x0800
#define ET_ARP  0x0806

struct arp_frame {
  __u16 hwtype;
  __u16 ptype;
  __u8  hlen;
  __u8  plen;  
  __u16 op;
  __u8  sha[6];
  __u8  spa[4];
  __u8  tha[6];
  __u8  tpa[4];
};

#define ARP_OP_REQUEST  1
#define ARP_OP_RESPONSE 2

#define ARP_HW_ETHERNET 1

#define HWADDR_LEN 6
#define IPADDR_LEN 4

// --------------------------------------------------------------------------------

struct hwif {
  char *name;
  struct sockaddr_ll *hwaddr;
  struct sockaddr_in *ipaddr;
  struct sockaddr_in *ipmask;
};

typedef struct hwif hwif_t;
typedef __u32 ip_t;

// ================================================================================

static int parse_ip(char *s, ip_t *ip) {
  int i;
  char *group;

  *ip = 0;

  for (i = 3,
       group = strtok(s, ". ");
       i >= 0 && group; 
       --i,
       group = strtok(NULL, ". ")) {
    *ip |= atoi(group) << (i*8);
  }
 
  return i == -1;
}

// --------------------------------------------------------------------------------

static void usage() {
  printf("Usage: arp <IP ADDRESS>\n");
}

// --------------------------------------------------------------------------------

#define HWIF_NUM_MAX 16

static struct hwif ifs[HWIF_NUM_MAX];
static int hwif_num = 0;

static hwif_t *lookup_if(char *name) {
  for (int i = 0; i < hwif_num; ++i) {
    if (!strcmp(ifs[i].name, name)) {
      return &ifs[i];
    }
  }

  // Create a new entry for the hardware interface
  if (hwif_num < HWIF_NUM_MAX) {
    ifs[hwif_num].name = name;
    ifs[hwif_num].hwaddr = NULL;
    ifs[hwif_num].ipaddr = NULL;
    return &ifs[hwif_num++];
  } else {
    return NULL;
  }
}

// --------------------------------------------------------------------------------

static struct ifaddrs *collect_hwif() {
  struct ifaddrs *ifa;
  struct ifaddrs *iface;

  getifaddrs(&ifa);

  /* 
     We have to find AF_INET and AF_PACKET famile addresses corresponding
     to a single interface.
  */
  for (iface = ifa;
       iface;
       iface = iface->ifa_next) {
    hwif_t *hwifp = lookup_if(iface->ifa_name);

    if (hwifp == NULL) {
      printf("Too many network interfaces. Using first %d of them.\n", HWIF_NUM_MAX);
      break;
    }
    
    if (iface->ifa_addr->sa_family == AF_PACKET) {  // Filter an Ethernet address
      struct sockaddr_ll *sap = (struct sockaddr_ll*)(iface->ifa_addr);

      if (sap->sll_hatype == ARPHRD_ETHER) {
        printf("Found an Ethernet interface: %s\n", iface->ifa_name);

        hwifp->hwaddr = sap;
      }
    } else if (iface->ifa_addr->sa_family == AF_INET) {  //Filter an IPv4 address
      hwifp->ipaddr = (struct sockaddr_in*)(iface->ifa_addr);
      hwifp->ipmask = (struct sockaddr_in*)(iface->ifa_netmask);

      printf("Found an IP interface: %s\n", iface->ifa_name);
    }
  }

  return ifa;
}

static int send_arp_request(hwif_t *ifp, int sock, ip_t ip) {
  struct sockaddr_ll sa;
  struct arp_frame af;

  sa.sll_family = AF_PACKET;
  sa.sll_protocol = htons(ETH_P_ARP);
  sa.sll_ifindex = ifp->hwaddr->sll_ifindex;
  sa.sll_hatype = ARPHRD_ETHER;
  sa.sll_pkttype = PACKET_BROADCAST;
  sa.sll_halen = HWADDR_LEN;
  memset(sa.sll_addr, 0xFF, HWADDR_LEN);
    
  af.hwtype = htons(ARP_HW_ETHERNET);
  af.ptype = htons(ET_IPV4);
  af.plen = IPADDR_LEN;
  af.hlen = HWADDR_LEN;
  af.op = htons(ARP_OP_REQUEST);
  memcpy(af.sha, ifp->hwaddr->sll_addr, HWADDR_LEN);
  memcpy(af.spa, &ifp->ipaddr->sin_addr, IPADDR_LEN);
  memcpy(af.tpa, &ip, IPADDR_LEN);
        
  printf("Sending an ARP request via %s\n", ifp->name);

  return sendto(sock, &af, sizeof(af), 0, (struct sockaddr*)&sa, sizeof(sa));  
}


#define FATAL -1
#define AGAIN 0
#define OK    1

static int timeout = 2000;

static int receive_arp_response(int sock, ip_t ip) {
  struct arp_frame af;
  struct pollfd pfd;
  int ret;  

  pfd.fd = sock;
  pfd.events = POLLIN;

  ret = poll(&pfd, 1, timeout);
  if (ret == 0) {
    printf("Waiting for an ARP response timed out.\n");

    return ret;
  } else if (ret < 0) {
    printf("Error waiting for an ARP response: %s\n", strerror(errno));

    return ret;
  }

  ret = recv(sock, &af, sizeof(af), 0);
  if (ret == sizeof(af)) {
    if (af.op == htons(ARP_OP_RESPONSE) && memcmp(af.spa, &ip, IPADDR_LEN) == 0) {
      printf("%d.%d.%d.%d is at %02X:%02X:%02X:%02X:%02X:%02X\n", 
             af.spa[0], af.spa[1], af.spa[2], af.spa[3], 
             af.sha[0], af.sha[1], af.sha[2], af.sha[3], af.sha[4], af.sha[5]);

      ret = OK;
    } else {
      printf("An ARP response from a non-target host received.\n");

      ret = AGAIN;
    }
  } else if (ret > 0) {
    printf("A malformed ARP response received.\n");
    ret = AGAIN;
  }

  return ret;
}


int main(int argc, char *argv[]) {
  struct ifaddrs *ifa;
  ip_t ip;
  int i;
  int ret;

  if (argc < 2 || !parse_ip(argv[1], &ip)) {
    usage();
    return 1;
  }

  ip = htonl(ip);
  ifa = collect_hwif();

  for (i = 0; i < hwif_num; ++i) {
    if (ifs[i].ipaddr && ifs[i].hwaddr) {
      ip_t if_ip, if_mask;

      if_ip = ifs[i].ipaddr->sin_addr.s_addr;
      if_mask = ifs[i].ipmask->sin_addr.s_addr;

      if ((if_ip & if_mask) == (ip & if_mask)) {
        break;
      }
    }
  }

  ret = 1;

  if (i == hwif_num) {
    printf("The host is unreachable.\n");

    goto cleanup;
  }

  int sock = socket(AF_PACKET, SOCK_DGRAM, ETH_P_ARP);
  if (sock == -1) {
    printf("Unable to create a link-level socket: %s\n", strerror(errno));
    goto cleanup;
  }

  int retries = 10;
  
  ifs[i].hwaddr->sll_protocol = htons(ETH_P_ARP);
  if (bind(sock, (struct sockaddr*)ifs[i].hwaddr, sizeof(struct sockaddr_ll))) {
    printf("Failed to bind to the interface %s!: %s\n", ifs[i].name, strerror(errno));
    
    close(sock);
    goto cleanup;
  }

  for (; retries > 0; retries--) {
    if (send_arp_request(&ifs[i], sock, ip) != sizeof(struct arp_frame)) {
      printf("Failed to send an ARP message via interface %s: %s\n", ifs[i].name, strerror(errno));
      break;
    }

    int res = receive_arp_response(sock, ip);
    if (res != AGAIN) {
      break;
    }
  }

  if (retries == 0) {
    printf("The target hasn't responded.");
  }
  
  close(sock);

  ret = 0;

 cleanup:  
  freeifaddrs(ifa);
  return ret;
}
